"""This script performs the FFT on a set of data and returns plots of the pre, during and post stimulus fourier plots, In addition a PSD can be plotted in a spectrogram format"""

from math import log
import numpy
import scipy
import pylab
from scipy import fftpack

def extract_data(filename):
    infile = open(filename, 'r')
    voltage = [] 
    times = []
    for line in infile:
        words = line.split() #this breaks each line into words
        volt = float(words[0])
        time = float(words[3])
        times.append(time)
        voltage.append(volt)
    infile.close()
    return voltage, times

voltage, times = extract_data('/home/davidmorton/Desktop/Data/043010RadialBar/273-Turtle_NI_Cortex_II-data') 
#pylab.plot(times,voltage)
#pylab.show()
pylab.suptitle('Power Spectrum Analysis 043010-273',fontsize=20)

##### Pre-stimulus Power Spectrum #################

v1 = numpy.array(voltage[0:49000])
t1 = numpy.array(times[0:49000])

Y1 = scipy.fftpack.fft(v1)
n1 = len(v1)
power1 = abs(Y1[1:(n1/2)])**2
nyquist = (1./.0001)/2
freq1 = numpy.array(range(n1/2))/(n1/2.0)*nyquist
#pylab.subplot(211)
pylab.subplot(231)
pylab.plot(freq1[1:len(freq1)],power1,label = 'Pre-stimulus')
pylab.ylabel('|FFT|$^2$', fontsize = 18)
pylab.axis([0, 50, 0, 5000])


##### During Stimulus Power spectra ################

v2 = numpy.array(voltage[50000:70000])
t2 = numpy.array(times[50000:70000])

Y2 = scipy.fftpack.fft(v2)
n2 = len(v2)
power2 = abs(Y2[1:(n2/2)])**2
nyquist = (1./.0001)/2
freq2 = numpy.array(range(n2/2))/(n2/2.0)*nyquist

pylab.subplot(232)
pylab.plot(freq2[1:len(freq2)],power2, label = 'During Stimulus')
pylab.axis([0, 50, 0, 5000])
pylab.setp(pylab.gca(),'yticklabels',[])
######## Post Stimulus Power Spectra ################

v3 = numpy.array(voltage[70000:120000])
t3 = numpy.array(times[70000:120000])

Y3 = scipy.fftpack.fft(v3)
n3 = len(v3)
power3 = abs(Y3[1:(n3/2)])**2
nyquist = (1./.0001)/2
freq3 = numpy.array(range(n3/2))/(n3/2.0)*nyquist

pylab.subplot(233)
pylab.plot(freq3[1:len(freq3)],power3)#label = 'Post-Stimulus')
pylab.axis([0, 50,0,5000])
pylab.setp(pylab.gca(),'yticklabels',[])




###################### Spectrogram ###############################
pylab.subplot(2,1,2)
Pxx, freqs, bins, im = pylab.specgram(voltage, NFFT = 8192, Fs = 1/0.0001)
pylab.ylim(0, 50)
pylab.xlabel('Time (s)', fontsize = 18)
pylab.ylabel('Frequency (Hz)',fontsize = 18)
pylab.show()

